#include "file_cmd.h"
#include <wchar.h>

BOOL set_pe_file_content(pe_file *p_pe, unsigned char *content, size_t len)
{
	p_pe->p_base = content;
	p_pe->length = len;
	p_pe->p_headers32 = (image_all_headers32 *)(p_pe->p_base + p_pe->p_dos_header->e_lfanew);
	
	p_pe->p_signature = &p_pe->p_headers32->Signature;
	p_pe->p_file_hdr = &p_pe->p_headers32->FileHeader;

	if(IS_64BIT_PE(p_pe))
	{
		p_pe->p_data_dir = p_pe->p_headers64->OptionalHeader.DataDirectory;
		p_pe->p_sect_hdr = p_pe->p_headers64->SectionHeaders;
	}
	else
	{
		p_pe->p_data_dir = p_pe->p_headers32->OptionalHeader.DataDirectory;
		p_pe->p_sect_hdr = p_pe->p_headers32->SectionHeaders;
	}

	if(!check_pe_headers(p_pe))
		return FALSE;

	if(!set_metadata(p_pe))
		return FALSE;

	return TRUE;
}

BOOL check_pe_headers(pe_file *p_pe)
{
	int n_dd_entries;

	if(p_pe->p_dos_header->e_magic != IMAGE_DOS_SIGNATURE)
	{
		wprintf(L"Missing DOS signature.\n");
		return FALSE;
	}

	if(p_pe->p_headers32->Signature != IMAGE_NT_SIGNATURE)
	{
		wprintf(L"Missing PE signature.\n");
		return FALSE;
	}

	switch(p_pe->p_headers32->OptionalHeader.Magic)
	{
	case IMAGE_NT_OPTIONAL_HDR32_MAGIC:
	case IMAGE_NT_OPTIONAL_HDR64_MAGIC:
		break;
	default:
		wprintf(L"Unknown optional header magic.\n");
		return FALSE;
	}

	n_dd_entries = IS_64BIT_PE(p_pe)
		? p_pe->p_headers64->OptionalHeader.NumberOfRvaAndSizes
		: p_pe->p_headers32->OptionalHeader.NumberOfRvaAndSizes;

	if(n_dd_entries != IMAGE_NUMBEROF_DIRECTORY_ENTRIES)
	{
		wprintf(L"Invalid number of data directory entries (%d), %d are required.\n",
			n_dd_entries, IMAGE_NUMBEROF_DIRECTORY_ENTRIES);
		return FALSE;
	}

	return TRUE;
}

void get_rva_info(pe_file *p_pe, u4 rva,
	image_section_header **pp_sec, size_t *p_sec_offset, size_t *p_file_offset)
{
	int i;

	for(i = 0; i < p_pe->p_file_hdr->NumberOfSections; ++i)
	{
		if(rva >= p_pe->p_sect_hdr[i].VirtualAddress
			&& rva < p_pe->p_sect_hdr[i].VirtualAddress + p_pe->p_sect_hdr[i].VirtualSize)
		{
			if(pp_sec != NULL)
				*pp_sec = &p_pe->p_sect_hdr[i];

			if(p_sec_offset != NULL)
				*p_sec_offset = rva - p_pe->p_sect_hdr[i].VirtualAddress;

			if(p_file_offset != NULL)
				*p_file_offset = rva - p_pe->p_sect_hdr[i].VirtualAddress
					+ p_pe->p_sect_hdr[i].PointerToRawData;

			return;
		}
	}

	// Reach here means not found
	if(pp_sec != NULL)
		*pp_sec = NULL;

	if(p_sec_offset != NULL)
		*p_sec_offset = 0;

	if(p_file_offset != NULL)
		*p_file_offset = 0;
}

void get_offset_info(pe_file *p_pe, u4 offset,
	image_section_header **pp_sec, size_t *p_sec_offset, size_t *p_rva)
{
	int i;

	for(i = 0; i < p_pe->p_file_hdr->NumberOfSections; ++i)
	{
		if(offset >= p_pe->p_sect_hdr[i].PointerToRawData
			&& offset < p_pe->p_sect_hdr[i].PointerToRawData + p_pe->p_sect_hdr[i].SizeOfRawData)
		{
			if(pp_sec != NULL)
				*pp_sec = &p_pe->p_sect_hdr[i];

			if(p_sec_offset != NULL)
				*p_sec_offset = offset - p_pe->p_sect_hdr[i].PointerToRawData;

			if(p_rva != NULL)
				*p_rva = offset - p_pe->p_sect_hdr[i].PointerToRawData
					+ p_pe->p_sect_hdr[i].VirtualAddress;

			return;
		}
	}

	// Reach here means not found
	if(pp_sec != NULL)
		*pp_sec = NULL;

	if(p_sec_offset != NULL)
		*p_sec_offset = 0;

	if(p_rva != NULL)
		*p_rva = 0;
}

size_t rva_to_offset(pe_file *p_pe, u4 rva)
{
	int i;
	image_section_header *p_sec = IS_64BIT_PE(p_pe)
		? p_pe->p_headers64->SectionHeaders
		: p_pe->p_headers32->SectionHeaders;

	for(i = 0; i < p_pe->p_headers32->FileHeader.NumberOfSections; ++i)
	{
		if(rva >= p_sec[i].VirtualAddress
			&& rva < p_sec[i].VirtualAddress + p_sec[i].VirtualSize)
		{
			return rva - p_sec[i].VirtualAddress + p_sec[i].PointerToRawData;
		}
	}

	return 0;
}

byte *round_to(byte *pointer, size_t alignment)
{
	size_t remain = (size_t)pointer % alignment;
	if(remain)
		pointer += alignment - remain;

	return pointer;
}

size_t format_rva_str(pe_file *p_pe, u4 rva, str_w str_buf, size_t buf_size)
{
	size_t len, sec_offset, file_offset;
	image_section_header *p_sh;

	get_rva_info(p_pe, rva, &p_sh, &sec_offset, &file_offset);

	if(p_sh)
	{
		for(len = 0;
			len < buf_size && len < IMAGE_SIZEOF_SHORT_NAME && p_sh->Name[len];
			++len)
			str_buf[len] = p_sh->Name[len];

		len += swprintf(str_buf + len, buf_size - len, L"+%Xh, FileOffset=%Xh",
			sec_offset, file_offset);
	}
	else
		len = swprintf(str_buf, buf_size, L"<invalid RVA>");

	return len;
}

size_t format_data_dir(pe_file *p_pe, image_data_directory *ddir,
	str_w str_buf, size_t buf_size)
{
	size_t len;

	len = swprintf(str_buf, buf_size, L"[%08Xh : %Xh (%d)]",
		ddir->VirtualAddress, ddir->Size, ddir->Size);

	if(ddir->VirtualAddress)
	{
		len += swprintf(str_buf + len, buf_size - len, L" (");
		len += format_rva_str(p_pe, ddir->VirtualAddress, str_buf + len, buf_size - len);
		len += swprintf(str_buf + len, buf_size - len, L")");
	}

	return len;
}

void exec_rva(pe_file *p_pe, size_t rva)
{
	image_section_header *p_sec;
	size_t sec_off, file_off;

	if(!check_file_opened() || !check_pe_headers(p_pe))
		return;

	get_rva_info(p_pe, rva, &p_sec, &sec_off, &file_off);

	if(p_sec == NULL)
	{
		wprintf(L"%08Xh is not a valid RVA value.\n", rva);
	}
	else
	{
		int i;
		wchar_t sec_name[IMAGE_SIZEOF_SHORT_NAME + 1];
		sec_name[IMAGE_SIZEOF_SHORT_NAME] = L'\0';
		for(i = 0; i < IMAGE_SIZEOF_SHORT_NAME; ++i)
			sec_name[i] = p_sec->Name[i];

		wprintf(L"           RVA: %08Xh\n", rva);
		wprintf(L"  Section Name: %s\n", sec_name);
		wprintf(L" Section Range: [%08Xh:%Xh (%d)]\n",
			p_sec->VirtualAddress, p_sec->VirtualSize, p_sec->VirtualSize);
		wprintf(L"Section Offset: %08Xh\n", sec_off);
		wprintf(L"   File Offset: %08Xh\n", file_off);
	}
}

void exec_offset(pe_file *p_pe, size_t offset)
{
	image_section_header *p_sec;
	size_t sec_off, rva;

	if(!check_file_opened() || !check_pe_headers(p_pe))
		return;

	get_offset_info(p_pe, offset, &p_sec, &sec_off, &rva);

	if(p_sec == NULL)
	{
		wprintf(L"%08Xh belongs to no section.\n", offset);
	}
	else
	{
		int i;
		wchar_t sec_name[IMAGE_SIZEOF_SHORT_NAME + 1];
		sec_name[IMAGE_SIZEOF_SHORT_NAME] = L'\0';
		for(i = 0; i < IMAGE_SIZEOF_SHORT_NAME; ++i)
			sec_name[i] = p_sec->Name[i];

		wprintf(L"        Offset: %08Xh\n", offset);
		wprintf(L"  Section Name: %s\n", sec_name);
		wprintf(L" Section Range: [%08Xh:%Xh (%d)]\n",
			p_sec->VirtualAddress, p_sec->VirtualSize, p_sec->VirtualSize);
		wprintf(L"Section Offset: %08Xh\n", sec_off);
		wprintf(L"           RVA: %08Xh\n", rva);
	}
}
